Add dataset type olmo_grain for AI2 OLMo numpy pretrain mixes#3749
Add dataset type olmo_grain for AI2 OLMo numpy pretrain mixes#3749
Conversation
|
🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details. |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
gagika
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
dipannita08,
gobbleturk,
hengtaoguo,
igorts-git,
jacoguzo,
jesselu-google,
jiangjy1982,
khatwanimohit,
richjames0,
shralex,
suexu1025 and
vipannalla
as code owners
|
🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
This Pull Request introduces a high-quality, Grain-based input pipeline for AI2's OLMo numpy datasets. The implementation is robust, well-documented, and includes a particularly clean approach to stateless resumption by deriving the data offset from the model checkpoint step.
🔍 General Feedback
- Stateless Resume: The
initial_steplogic in the sampler is an excellent design choice that avoids the complexities of Grain iterator-state serialization. - N-gram Filtering: The integration of OLMo-core's repetition filter via a custom transform that masks instances in the loss is both efficient and sharding-friendly.
- Testing and Validation: The inclusion of unit tests, smoke scripts, and end-to-end resume tests provides great confidence in the stability of the new pipeline.
- Performance: While the in-memory permutation for shuffling is currently manageable, it's worth monitoring as dataset sizes scale further.
| reverse = _find_end_first_consecutive_true(arr[::-1]) | ||
| return len(arr) - reverse if reverse > 0 else -1 | ||
|
|
||
|
|
There was a problem hiding this comment.
🟡 The slicing prog[:-1:] is functional but non-idiomatic in Python. Using prog[:-1] is cleaner and more conventional.
| true_locs = np.where(prog[:-1] == prog[1:])[0] |
|
|
||
| max_period = min(max_period, len(arr) // 3) | ||
|
|
||
| for period in range(min_period, max_period + 1): |
There was a problem hiding this comment.
🟢 Using a mask_value of -1 to pad rows before reshaping is a clever way to avoid periodic matches wrapping around rows. However, if -1 (or its uint32 representation 0xFFFFFFFF) is a valid token in the dataset, the ValueError will be triggered unnecessarily.
Consider using a value that is strictly outside the tokenizer's vocabulary range if available, or documenting that -1 is reserved.
| total_instances: ``index.total_instances`` from the OLMo index. | ||
| seed: Base seed for the shuffle. | ||
| shard_index: Zero-based index of this data-loading host. Typically | ||
| ``jax.process_index()``. |
There was a problem hiding this comment.
🟡 For very large datasets (e.g., the 724M instance mix mentioned), allocating the full permutation in host memory (~5.8 GB) can be a significant spike, especially if many hosts are doing it simultaneously at an epoch boundary.
While acceptable for the current scope, consider implementing a lazy or on-disk permutation scheme if the dataset size grows further or if host memory becomes a constraint.
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
Thanks for the change!
Have you had a chance to chat with @aireenmei about this yet? I'm wondering if we could design these features to directly leverage the existing grain_data_processing.py. For example, we could add things like n-gram filtering and pre-tokenized numpy mixes as features there to improve code reusability.
The main benefit I see is reducing maintenance overhead. We currently maintain both tfds and c4_tfds_mlperf, but the latter is rarely used and has some maintenance issues. Since Grain will be heavily used moving forward, it makes sense to build on top of it directly. Let me know your thoughts—happy to discuss!
|
|
||
| ```bash | ||
| python tools/data_generation/download_olmo_data_to_gcs.py \ | ||
| --mix-file /path/to/OLMo-mix-0925-official.txt \ |
There was a problem hiding this comment.
Will this path be at GCS or not necessary?
| """ | ||
|
|
||
| olmo_index_path: PathStr = Field("", description="Path or gs:// URI to the JSON index from build_olmo_npy_index.py.") | ||
| olmo_path_remap_from: PathStr = Field( |
There was a problem hiding this comment.
Maybe add some comments and indicate why we need this?
3 participants
Description
dataset_type=olmo_grain, a Grain-based input pipeline for AI2'spre-tokenized OLMo numpy mixes (e.g.
OLMo-mix-0925-official.txt).Reads headerless
.npytoken streams from a gcsfuse mount, appliesOLMo-core's repeated-n-gram filter, and yields the shapes the MaxText
pretrain trainer expects.
(seed, shard, k). Resume reads the latest step fromconfig.checkpoint_dirand shifts the sampler — no Grain iteratorstate in the checkpoint.
download_olmo_data_to_gcs.pywith HTTP-Rangeresume;
build_olmo_npy_index.pyfor header-scan indexing) and twolaunchers (
run_olmo3_7b_grain_smoke.sh,run_olmo3_7b_grain_resume_test.sh).Tests
tests/unit/input_pipeline/olmo_*)Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.authored-by: @aireenmei